#!/bin/bash
###############################################################################
#       脚本名称: filter_rate_limit.sh
#       功能描述: 频率限制过滤器,防止日志洪水攻击
#       作者: 林再来
#       日期: 2025-10-14
#       版本: 1.0
#       功能:
#           1、按时间窗口限制日志数量
#           2、按消息内容去重
#           3、滑动窗口计数
#           4、自动清理过期记录
#       依赖(插件/命令):
#           无
###############################################################################

# 防止重复加载
[[ -n "${_FILTER_RATE_LIMIT_LOADED:-}" ]] && return 0
readonly _FILTER_RATE_LIMIT_LOADED=1

#-----------------------------------------------------------------------------
# 全局变量
#-----------------------------------------------------------------------------
# 频率限制配置
declare -gA _FILTER_RATE_LIMIT_CONFIG
# 计数器(时间戳列表)
declare -gA _FILTER_RATE_LIMIT_COUNTER
# 消息哈希记录(用于去重)
declare -gA _FILTER_RATE_LIMIT_HASH

#-----------------------------------------------------------------------------
# 默认配置（参考config模块的做法，先检查是否已定义）
#-----------------------------------------------------------------------------
if [[ -z "${DEFAULT_MAX_LOGS_PER_SECOND:-}" ]]; then
    readonly DEFAULT_MAX_LOGS_PER_SECOND=100
fi
if [[ -z "${DEFAULT_TIME_WINDOW:-}" ]]; then
    readonly DEFAULT_TIME_WINDOW=60          # 60秒时间窗口
fi
if [[ -z "${DEFAULT_ENABLE_DEDUP:-}" ]]; then
    readonly DEFAULT_ENABLE_DEDUP="false"    # 默认不启用去重
fi

#-----------------------------------------------------------------------------
# 函数: 初始化频率限制过滤器
# 参数: $1 - 过滤器名称
#       $2 - 最大日志数/秒(可选,默认100)
#       $3 - 时间窗口/秒(可选,默认60)
#       $4 - 是否启用去重(可选,默认false)
# 返回: 0-成功
#-----------------------------------------------------------------------------
function _filter_rate_limit_init() {
    local filter_name="$1"
    local max_logs="${2:-$DEFAULT_MAX_LOGS_PER_SECOND}"
    local time_window="${3:-$DEFAULT_TIME_WINDOW}"
    local enable_dedup="${4:-$DEFAULT_ENABLE_DEDUP}"
    
    _FILTER_RATE_LIMIT_CONFIG["${filter_name}_max"]="$max_logs"
    _FILTER_RATE_LIMIT_CONFIG["${filter_name}_window"]="$time_window"
    _FILTER_RATE_LIMIT_CONFIG["${filter_name}_dedup"]="$enable_dedup"
    _FILTER_RATE_LIMIT_COUNTER["$filter_name"]=""
    _FILTER_RATE_LIMIT_HASH["$filter_name"]=""
    
    return 0
}

#-----------------------------------------------------------------------------
# 函数: 计算字符串哈希(简单实现)
# 参数: $1 - 字符串
# 返回: 输出哈希值
#-----------------------------------------------------------------------------
function _compute_hash() {
    local str="$1"
    # 使用cksum计算简单哈希
    echo "$str" | cksum | awk '{print $1}'
}

#-----------------------------------------------------------------------------
# 函数: 清理过期记录
# 参数: $1 - 过滤器名称
# 返回: 0-成功
#-----------------------------------------------------------------------------
function _filter_rate_limit_cleanup() {
    local filter_name="$1"
    local current_time=$(date +%s)
    local time_window="${_FILTER_RATE_LIMIT_CONFIG[${filter_name}_window]}"
    local cutoff_time=$((current_time - time_window))
    
    # 清理计数器中的过期时间戳
    local timestamps="${_FILTER_RATE_LIMIT_COUNTER[$filter_name]}"
    local new_timestamps=""
    
    if [[ -n "$timestamps" ]]; then
        for ts in $timestamps; do
            if [[ $ts -gt $cutoff_time ]]; then
                new_timestamps+="$ts "
            fi
        done
    fi
    
    _FILTER_RATE_LIMIT_COUNTER["$filter_name"]="${new_timestamps% }"
    
    # 清理哈希表(简单实现:全部清空)
    # TODO: 更精细的清理策略
    if [[ -n "${_FILTER_RATE_LIMIT_HASH[$filter_name]}" ]]; then
        _FILTER_RATE_LIMIT_HASH["$filter_name"]=""
    fi
    
    return 0
}

#-----------------------------------------------------------------------------
# 函数: 应用频率限制过滤
# 参数: $1 - 过滤器名称
#       $2 - 日志事件JSON
# 返回: 0-通过, 1-拒绝(超出限制)
#-----------------------------------------------------------------------------
function _filter_rate_limit_apply() {
    local filter_name="$1"
    local event="$2"
    
    local current_time=$(date +%s)
    local max_logs="${_FILTER_RATE_LIMIT_CONFIG[${filter_name}_max]}"
    local time_window="${_FILTER_RATE_LIMIT_CONFIG[${filter_name}_window]}"
    local enable_dedup="${_FILTER_RATE_LIMIT_CONFIG[${filter_name}_dedup]}"
    
    # 提取消息内容
    local message=$(echo "$event" | grep -oP '(?<="message":")[^"]*' 2>/dev/null)
    
    # 去重检查
    if [[ "$enable_dedup" == "true" ]]; then
        local msg_hash=$(_compute_hash "$message")
        local hash_list="${_FILTER_RATE_LIMIT_HASH[$filter_name]}"
        
        if [[ "$hash_list" == *"$msg_hash"* ]]; then
            return 1  # 重复消息,拒绝
        fi
        
        # 添加到哈希列表
        _FILTER_RATE_LIMIT_HASH["$filter_name"]+="$msg_hash "
    fi
    
    # 清理过期记录(每10秒清理一次)
    if [[ $((current_time % 10)) -eq 0 ]]; then
        _filter_rate_limit_cleanup "$filter_name"
    fi
    
    # 统计时间窗口内的日志数量
    local cutoff_time=$((current_time - time_window))
    local timestamps="${_FILTER_RATE_LIMIT_COUNTER[$filter_name]}"
    local count=0
    
    if [[ -n "$timestamps" ]]; then
        for ts in $timestamps; do
            if [[ $ts -gt $cutoff_time ]]; then
                ((count++))
            fi
        done
    fi
    
    # 计算每秒速率
    local rate=$((count / time_window))
    
    # 检查是否超出限制
    if [[ $rate -ge $max_logs ]]; then
        return 1  # 超出限制,拒绝
    fi
    
    # 记录当前时间戳
    _FILTER_RATE_LIMIT_COUNTER["$filter_name"]+="$current_time "
    
    return 0  # 通过
}

#-----------------------------------------------------------------------------
# 函数: 获取当前速率
# 参数: $1 - 过滤器名称
# 返回: 输出当前日志速率(条/秒)
#-----------------------------------------------------------------------------
function _filter_rate_limit_get_rate() {
    local filter_name="$1"
    local current_time=$(date +%s)
    local time_window="${_FILTER_RATE_LIMIT_CONFIG[${filter_name}_window]}"
    local cutoff_time=$((current_time - time_window))
    local timestamps="${_FILTER_RATE_LIMIT_COUNTER[$filter_name]}"
    local count=0
    
    if [[ -n "$timestamps" ]]; then
        for ts in $timestamps; do
            if [[ $ts -gt $cutoff_time ]]; then
                ((count++))
            fi
        done
    fi
    
    local rate=$((count / time_window))
    echo "$rate"
}

#-----------------------------------------------------------------------------
# 函数: 重置频率限制过滤器
# 参数: $1 - 过滤器名称
# 返回: 0-成功
#-----------------------------------------------------------------------------
function _filter_rate_limit_reset() {
    local filter_name="$1"
    
    _FILTER_RATE_LIMIT_COUNTER["$filter_name"]=""
    _FILTER_RATE_LIMIT_HASH["$filter_name"]=""
    
    return 0
}

#-----------------------------------------------------------------------------
# 函数: 获取统计信息
# 参数: $1 - 过滤器名称
# 返回: 输出统计信息
#-----------------------------------------------------------------------------
function _filter_rate_limit_stats() {
    local filter_name="$1"
    
    local max_logs="${_FILTER_RATE_LIMIT_CONFIG[${filter_name}_max]}"
    local time_window="${_FILTER_RATE_LIMIT_CONFIG[${filter_name}_window]}"
    local enable_dedup="${_FILTER_RATE_LIMIT_CONFIG[${filter_name}_dedup]}"
    local current_rate=$(_filter_rate_limit_get_rate "$filter_name")
    
    echo "Rate Limit Filter: $filter_name"
    echo "  最大速率: ${max_logs}条/秒"
    echo "  时间窗口: ${time_window}秒"
    echo "  启用去重: $enable_dedup"
    echo "  当前速率: ${current_rate}条/秒"
}

#-----------------------------------------------------------------------------
# 测试函数
#-----------------------------------------------------------------------------
function _filter_rate_limit_test() {
    echo "=== 频率限制过滤器测试 ==="
    
    echo ""
    echo "1. 初始化过滤器(最大5条/秒,10秒窗口):"
    _filter_rate_limit_init "test_filter" 5 10 "false"
    echo "  ✓ 初始化成功"
    
    echo ""
    echo "2. 快速发送日志测试:"
    local passed=0
    local rejected=0
    
    for i in {1..20}; do
        local test_event="{\"message\":\"测试日志 $i\"}"
        if _filter_rate_limit_apply "test_filter" "$test_event"; then
            ((passed++))
        else
            ((rejected++))
        fi
    done
    
    echo "  通过: $passed 条"
    echo "  拒绝: $rejected 条"
    
    echo ""
    echo "3. 查看统计信息:"
    _filter_rate_limit_stats "test_filter" | sed 's/^/  /'
    
    echo ""
    echo "4. 测试去重功能:"
    _filter_rate_limit_reset "test_filter"
    _filter_rate_limit_init "test_filter" 100 10 "true"
    
    local test_event='{"message":"重复消息"}'
    echo "  第1次发送: $(_filter_rate_limit_apply "test_filter" "$test_event" && echo 通过 || echo 拒绝)"
    echo "  第2次发送(相同消息): $(_filter_rate_limit_apply "test_filter" "$test_event" && echo 通过 || echo 拒绝)"
    
    local test_event2='{"message":"不同消息"}'
    echo "  第3次发送(不同消息): $(_filter_rate_limit_apply "test_filter" "$test_event2" && echo 通过 || echo 拒绝)"
}

# 如果直接执行此脚本,运行测试
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    _filter_rate_limit_test
fi

